#pragma once

#include <iostream>
#include <list>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>

#include "arithmetic.h"
#include "effect.h"
#include "iroperand.h"
#include "symbol.h"
#include "utils.h"

struct IRBlock;
struct IRFunction;

namespace IR {
    struct IRBase {
        IRBlock* parent = nullptr;

        virtual void Print() const = 0;
        virtual std::vector<Variable*> UsedVar() const = 0;
        virtual std::vector<Variable*> ChangedVar() const = 0;
        virtual SideEffect Effects() const = 0;

        template <class T>
        T* As() {
            return dynamic_cast<T*>(this);
        }
    };

    // Operation - Binary operation (lhs != nullptr) or unary operation
    // dest := lhs op rhs
    struct Operation : public IRBase {
        Variable* dest = nullptr;
        ArithOp op = ArithOp::Add;
        Operand* lhs = nullptr;
        Operand* rhs = nullptr;

        // dest := rhs
        bool IsCopy() const { return op == ArithOp::Add && !lhs; }
        bool IsNop() const { return !dest; }

        void Print() const override;
        std::vector<Variable*> UsedVar() const override;
        std::vector<Variable*> ChangedVar() const override;
        SideEffect Effects() const override;
    };

    // Alloc - dest := new (size)
    struct Alloc : public IRBase {
        Variable* dest;
        int size;

        void Print() const override;
        std::vector<Variable*> UsedVar() const override;
        std::vector<Variable*> ChangedVar() const override;
        SideEffect Effects() const override;
    };

    // Load - dest := base[offset]
    struct Load : public IRBase {
        Variable* dest;
        Variable* base;
        Operand* offset;

        void Print() const override;
        std::vector<Variable*> UsedVar() const override;
        std::vector<Variable*> ChangedVar() const override;
        SideEffect Effects() const override;
    };

    // Store - base[offset] := src
    struct Store : public IRBase {
        Variable* base;
        Operand* offset;
        Operand* src;

        void Print() const override;
        std::vector<Variable*> UsedVar() const override;
        std::vector<Variable*> ChangedVar() const override;
        SideEffect Effects() const override;
    };

    // Phi - dest := phi(srcs)
    struct Phi : public IRBase {
        Variable* dest;
        std::unordered_map<IRBlock*, Variable*> srcs;

        void Print() const override;
        std::vector<Variable*> UsedVar() const override;
        std::vector<Variable*> ChangedVar() const override;
        SideEffect Effects() const override;
    };

    // Call - dest := func(args)
    struct Call : public IRBase {
        Variable* dest;
        IRFunction* func;
        std::vector<Operand*> args;

        void Print() const override;
        std::vector<Variable*> UsedVar() const override;
        std::vector<Variable*> ChangedVar() const override;
        SideEffect Effects() const override;
    };

    // Terminator - jump, return, branch
    struct Terminator : public IRBase {
        virtual std::vector<IRBlock*> GetNext() const = 0;
    };

    // Return - return ret
    struct Return : public Terminator {
        Operand* ret;

        void Print() const override;
        std::vector<Variable*> UsedVar() const override;
        std::vector<Variable*> ChangedVar() const override;
        SideEffect Effects() const override;
        std::vector<IRBlock*> GetNext() const override;
    };

    // Jump - goto nxt
    struct Jump : public Terminator {
        IRBlock* nxt;

        void Print() const override;
        std::vector<Variable*> UsedVar() const override;
        std::vector<Variable*> ChangedVar() const override;
        SideEffect Effects() const override;
        std::vector<IRBlock*> GetNext() const override;
    };

    // Branch - goto (lhs op rhs) ? t : f
    struct Branch : public Terminator {
        ArithOp op = ArithOp::Add;
        Operand* lhs = nullptr;
        Operand* rhs = nullptr;
        IRBlock* t;
        IRBlock* f;

        void Print() const override;
        std::vector<Variable*> UsedVar() const override;
        std::vector<Variable*> ChangedVar() const override;
        SideEffect Effects() const override;
        std::vector<IRBlock*> GetNext() const override;
    };

}  // namespace IR
